require 'rake/clean'
require 'pp'
require 'yaml'
require 'securerandom'
require 'fileutils'
require 'beaker-hostgenerator'

$LOAD_PATH << File.expand_path(File.join(File.dirname(__FILE__), 'lib'))
require 'puppet/acceptance/git_utils'
extend Puppet::Acceptance::GitUtils

ONE_DAY_IN_SECS = 24 * 60 * 60
REPO_CONFIGS_DIR = "repo-configs"
CLEAN.include('*.tar', REPO_CONFIGS_DIR, 'merged_options.rb')

module HarnessOptions

  DEFAULTS = {
    :type => 'git',
    :helper => ['lib/helper.rb'],
    :tests  => ['tests'],
    :log_level => 'debug',
    :color => false,
    :root_keys => true,
    :ssh => {
      :keys => ["~/.ssh/id_rsa-acceptance"],
    },
    :xml => true,
    :timesync => false,
    :repo_proxy => true,
    :add_el_extras => false,
    :preserve_hosts => 'onfail',
  }

  class Aggregator
    attr_reader :mode

    def initialize(mode)
      @mode = mode
    end

    def get_options(file_path)
      puts file_path
      if File.exists? file_path
        options = eval(File.read(file_path), binding)
      else
        puts "No options file found at #{File.expand_path(file_path)}"
      end
      options || {}
    end

    def get_mode_options
      get_options("./config/#{mode}/options.rb")
    end

    def get_local_options
      get_options("./local_options.rb")
    end

    def final_options(intermediary_options = {})
      mode_options = get_mode_options
      local_overrides = get_local_options
      final_options = DEFAULTS.merge(mode_options)
      final_options.merge!(intermediary_options)
      final_options.merge!(local_overrides)
      return final_options
    end
  end

  def self.options(mode, options)
    final_options = Aggregator.new(mode).final_options(options)
    final_options
  end
end

def beaker_test(mode = :aio, options = {})
  final_options = HarnessOptions.options(mode, options)

  if mode == :git
    final_options[:install] ||= []
    final_options[:install] << "#{build_giturl('facter')}##{sha}"
  end

  options_file = 'merged_options.rb'
  File.open(options_file, 'w') do |merged|
    merged.puts <<-EOS
# Copy this file to local_options.rb and adjust as needed if you wish to run
# with some local overrides.
EOS
    merged.puts(final_options.pretty_inspect)
  end

  tests = ENV['TESTS'] || ENV['TEST']
  tests_opt = "--tests=#{tests}" if tests

  target = ENV['TEST_TARGET']
  if config and File.exists?(config)
    config_opt = "--hosts=#{config}"
  elsif target
    cli = BeakerHostGenerator::CLI.new([target, '--disable-default-role', '--osinfo-version', '1'])
    ENV['BEAKER_HOSTS'] = "tmp/#{target}-#{SecureRandom.uuid}.yaml"
    FileUtils.mkdir_p('tmp')
    File.open(config, 'w') do |fh|
      fh.print(cli.execute)
    end
    config_opt = "--hosts=#{config}"
  end

  overriding_options = ENV['OPTIONS']

  args = ["--options-file", options_file, config_opt, tests_opt, overriding_options].compact

  begin
    sh("beaker", *args)
  ensure
    if (hosts_file = config || final_options[:hosts_file]) && hosts_file !~ /preserved_config/
      cp(hosts_file, "log/latest/config.yml")
      generate_config_for_latest_hosts if final_options[:preserve_hosts] || overriding_options =~ /--preserve-hosts/
    end
    mv(options_file, "log/latest")
    puts "\n\n"
  end
end

def generate_config_for_latest_hosts
  preserved_config_hash = { 'HOSTS' => {} }

  config_hash = YAML.load_file('log/latest/config.yml').to_hash
  nodes = config_hash['HOSTS'].map do |node_label,hash|
    { :node_label => node_label, :platform => hash['platform'] }
  end

  pre_suite_log = File.read('log/latest/pre_suite-run.log')
  nodes.each do |node_info|
    possible_hostname_match = /^(\w+) \(#{node_info[:node_label]}\)/.match(pre_suite_log)
    hostname = (possible_hostname_match || Array.new)[1]
    fqdn = hostname ? "#{hostname}.delivery.puppetlabs.net" : "unknown"
    preserved_config_hash['HOSTS'][fqdn] = {
      'roles' => [ 'agent'],
      'platform' => node_info[:platform],
    }
    preserved_config_hash['HOSTS'][fqdn]['roles'].unshift('master') if node_info[:node_label] =~ /master/
  end

  puts "\n\n"
  pp preserved_config_hash
  puts "\n\n"

  File.open('log/latest/preserved_config.yaml', 'w') do |config_file|
    YAML.dump(preserved_config_hash, config_file)
  end
rescue Errno::ENOENT => e
  puts "Couldn't generate log #{e}"
end

def list_preserved_configurations(secs_ago = ONE_DAY_IN_SECS)
  preserved = {}
  Dir.glob('log/*_*').each do |dir|
    preserved_config_path = "#{dir}/preserved_config.yaml"
    yesterday = Time.now - secs_ago.to_i
    if preserved_config = File.exists?(preserved_config_path)
      directory = File.new(dir)
      if directory.ctime > yesterday
        hosts = []
        preserved_config = YAML.load_file(preserved_config_path).to_hash
        preserved_config['HOSTS'].each do |hostname,values|
          hosts << "#{hostname}: #{values['platform']}, #{values['roles']}"
        end
        preserved[hosts] = directory.to_path
      end
    end
  end
  preserved.map { |k,v| [v,k] }.sort { |a,b| a[0] <=> b[0] }.reverse
end

def list_preserved_hosts(secs_ago = ONE_DAY_IN_SECS)
  hosts = Set.new
  Dir.glob('log/**/pre*suite*run.log').each do |log|
    yesterday = Time.now - secs_ago.to_i
    File.open(log, 'r') do |file|
      if file.ctime > yesterday
        file.each_line do |line|
          matchdata = /^(\w+) \(.*?\) \$/.match(line.encode!('UTF-8', 'UTF-8', :invalid => :replace))
          hosts.add(matchdata[1]) if matchdata
        end
      end
    end
  end
  hosts
end

# Plagiarized from Beaker::Vcloud#cleanup
def destroy_preserved_hosts(hosts = nil, secs_ago = ONE_DAY_IN_SECS)
  secs_ago ||= ONE_DAY_IN_SECS
  hosts ||= list_preserved_hosts(secs_ago)

  require 'beaker/hypervisor/vsphere_helper'
  vsphere_credentials = VsphereHelper.load_config("#{ENV['HOME']}/.fog")

  puts "Connecting to vSphere at #{vsphere_credentials[:server]}" +
    " with credentials for #{vsphere_credentials[:user]}"

  vsphere_helper = VsphereHelper.new( vsphere_credentials )

  vm_names = hosts.to_a
  pp vm_names
  vms = vsphere_helper.find_vms vm_names
  vm_names.each do |name|
    unless vm = vms[name]
      puts "Couldn't find VM #{name} in vSphere!"
      next
    end

    if vm.runtime.powerState == 'poweredOn'
      puts "Shutting down #{vm.name}"
      start = Time.now
      vm.PowerOffVM_Task.wait_for_completion
      puts "Spent %.2f seconds halting #{vm.name}" % (Time.now - start)
    end

    start = Time.now
    vm.Destroy_Task
    puts "Spent %.2f seconds destroying #{vm.name}" % (Time.now - start)
  end

  vsphere_helper.close
end

def print_preserved(preserved)
  preserved.each_with_index do |entry,i|
    puts "##{i}: #{entry[0]}"
    entry[1].each { |h| puts "  #{h}" }
  end
end

def beaker_run_type
  type = ENV['TYPE'] || :aio
  type = type.to_sym
end

def sha
  ENV['SHA']
end

def config
  ENV['BEAKER_HOSTS']
end

namespace :ci do

  task :check_env do
    raise(USAGE) unless sha
  end

  namespace :test do

    USAGE = <<-EOS
Requires commit SHA to be put under test as environment variable: SHA='<sha>'.
Also must set BEAKER_HOSTS=config/nodes/foo.yaml or include it in an options.rb for Beaker,
or specify TEST_TARGET in a form beaker-hostgenerator accepts, e.g. ubuntu1504-64a.
You may set TESTS=path/to/test,and/more/tests.
You may set additional Beaker OPTIONS='--more --options'
If testing from git checkouts, you may optionally set the github fork to checkout from using FORK='other-puppet-fork'.
You may also optionally set the git server to checkout from using GIT_SERVER='my.host.with.git.daemon', if you have set up a `git daemon` to pull local commits from.  (In this case, the FORK should be set to the path to the repository, and you will need to allow the git daemon to serve the repo (see `git help daemon`)).
If there is a Beaker options hash in a ./local_options.rb, it will be included.  Commandline options set through the above environment variables will override settings in this file.
EOS

    desc <<-EOS
Run the acceptance tests through Beaker and install packages as part of the AIO puppet-agent installation.
#{USAGE}
EOS
    task :aio => 'ci:check_env' do
      beaker_test(:aio)
    end

    desc <<-EOS
Run the acceptance tests through Beaker and install from git on the configuration targets.
#{USAGE}
EOS
    task :git => 'ci:check_env' do
      beaker_test(:git)
    end
  end

  desc "Capture the master and agent hostname from the latest log and construct a preserved_config.yaml for re-running against preserved hosts without provisioning."
  task :extract_preserved_config do
    generate_config_for_latest_hosts
  end

  desc <<-EOS
Run an acceptance test for a given node configuration and preserve the hosts.
Defaults to a packages run, but you can set it to 'git' with TYPE='git'.
#{USAGE}
  EOS
  task :test_and_preserve_hosts => 'ci:check_env'  do
    beaker_test(beaker_run_type, :preserve_hosts => 'always', :__preserve_config__ => true)
  end

  desc "List acceptance runs from the past day which had hosts preserved."
  task :list_preserved do
    preserved = list_preserved_configurations
    print_preserved(preserved)
  end

  desc <<-EOS
Shutdown and destroy any hosts that we have preserved for testing.  These should be reaped daily by scripts, but this will free up resources immediately.
Specify a list of comma separated HOST_NAMES if you have a set of dynamic vcloud host names you want to purge outside of what can be grepped from the logs.
You can go back through the last SECS_AGO logs.  Default is one day ago in secs.
  EOS
  task :destroy_preserved_hosts do
    host_names = ENV['HOST_NAMES'].split(',') if ENV['HOST_NAMES']
    secs_ago = ENV['SECS_AGO']
    destroy_preserved_hosts(host_names, secs_ago)
  end

  desc <<-EOS
Rerun an acceptance test using the last captured preserved_config.yaml to skip provisioning.
Or specify a CONFIG_NUMBER from `rake ci:list_preserved`.
Uses the setup/rsync/pre-suite to rsync the local puppet source onto master and agent.
You may specify an RSYNC_FILTER_FILE as well.
You may skip purging and reinstalling puppet packages by including SKIP_PACKAGE_REINSTALL.
You may skip rsyncing local puppet files over to the tests hosts by including SKIP_RSYNC.
Defaults to an aio run, but you can set it to 'git' with TYPE='git'.
  EOS
  task :test_against_preserved_hosts do
    config_number = (ENV['CONFIG_NUMBER'] || 0).to_i
    preserved = list_preserved_configurations
    print_preserved(preserved)
    config_path = preserved[config_number][0]
    puts "Using ##{config_number}: #{config_path}"
    beaker_test(beaker_run_type,
      :hosts_file => "#{config_path}/preserved_config.yaml",
      :no_provision => true,
      :preserve_hosts => 'always',
      :pre_suite => ['setup/rsync/pre-suite']
    )
  end
end

task :default do
  sh('rake -T')
end
